package cmm.enhance;

import cmm.GlobalOptions;
import org.objectweb.asm.Opcodes;
import org.objectweb.asm.tree.*;

import java.util.Iterator;
import java.util.List;

import static org.objectweb.asm.Opcodes.*;
import static org.objectweb.asm.Opcodes.ATHROW;

/**
 * @author caomm
 * @Title 字节码增强 核心code
 * @Description TODO
 * @Date 2022/9/18 16:35
 */
public class MethodTraceTransform {

    public static void transform(ClassNode classNode, String classPath, String className, boolean isMainClass) {
        // class内methods
        List<MethodNode> methods = classNode.methods;

        for (MethodNode mn : methods) {
            InsnList methodInsn = mn.instructions;
            if ("<init>".equals(mn.name) || "<clinit>".equals(mn.name) || methodInsn.size() <= 0 || className.equals(mn.name)) {
                continue;
            }
            // 检查是否已经增强 : 未增强继续，否则跳过当前类
            if (checkAlreadyEnhance(mn)) {
                return;
            }

            // 前置增强
            int uuidKeyLocal = BeforeEnhance(mn, methodInsn, GlobalOptions.getClassMethod(classPath, mn.name));
            // 异常增强
            TryCatchEnhance(mn, methodInsn, isMainClass);
            // 后置增强
            AfterEnhance(methodInsn, GlobalOptions.getClassMethod(classPath, mn.name));
            // SQL耗时增强
            SqlEnhance(mn, methodInsn, GlobalOptions.getClassMethod(classPath, mn.name), uuidKeyLocal);
            // 外部调用增强
            HttpEnhance(mn, methodInsn, GlobalOptions.getClassMethod(classPath, mn.name), uuidKeyLocal);
        }
    }

    /**
     * 外部调用增强，目前只针对DAP框架外部调用invoke调用增强
     */
    private static void HttpEnhance(MethodNode mn, InsnList methodInsn, String methodPath, int uuidKeyLocal) {
        int maxLocals = mn.maxLocals;
        Iterator<AbstractInsnNode> j = methodInsn.iterator();
        while (j.hasNext()) {
            try {
                MethodInsnNode in = (MethodInsnNode) j.next();

                // 定位sql执行位置
                if ((in.owner.contains(GlobalOptions.HTTPCLIENTCLASS01) && GlobalOptions.isHttpClient01(in.name))
                        ||(in.owner.contains(GlobalOptions.HTTPCLIENTCLASS02) && GlobalOptions.isHttpClient02(in.name))) {
                    InsnList sqlBeforeInsn = new InsnList();
                    sqlBeforeInsn.add(new MethodInsnNode(Opcodes.INVOKESTATIC, "java/lang/System", "currentTimeMillis", "()J",
                            false));
                    sqlBeforeInsn.add(new VarInsnNode(LSTORE, maxLocals));

                    InsnList sqlAfterInsn = new InsnList();
                    sqlAfterInsn.add(new MethodInsnNode(Opcodes.INVOKESTATIC, "java/lang/System", "currentTimeMillis", "()J",
                            false));
                    sqlAfterInsn.add(new VarInsnNode(LLOAD, maxLocals));
                    sqlAfterInsn.add(new InsnNode(LSUB));
                    sqlAfterInsn.add(new LdcInsnNode(methodPath));
                    // uuidKeyLocal是在@before创建的uuid的局部变量表位置
                    sqlAfterInsn.add(new VarInsnNode(ALOAD, uuidKeyLocal));
                    sqlAfterInsn.add(new MethodInsnNode(Opcodes.INVOKESTATIC, "TraceMonitorHandler", "httpClientEnhance", "(JLjava/lang/String;Ljava/lang/String;)V",
                            false));

                    methodInsn.insert(in.getPrevious(), sqlBeforeInsn);
                    methodInsn.insert(in.getNext(), sqlAfterInsn);

                    mn.maxLocals += 1;
                    mn.maxStack += 1;
                }
            } catch (Exception e) {
                continue;
            }
        }
    }

    /**
     * SQL耗时增强 ，SQL执行位置定位目前针对DAP架构 com.digiwin.app.dao.DWDao 下的常用方法 例如：select、update、execute、delete
     */
    private static void SqlEnhance(MethodNode mn, InsnList methodInsn, String methodPath, int uuidKeyLocal) {
        int maxLocals = mn.maxLocals;
        Iterator<AbstractInsnNode> j = methodInsn.iterator();
        while (j.hasNext()) {
            try {
                MethodInsnNode in = (MethodInsnNode) j.next();

                // 定位sql执行位置
                if (in.owner.contains(GlobalOptions.SQLDAOCLASS) && GlobalOptions.isSqlExecute(in.name)) {
                    InsnList sqlBeforeInsn = new InsnList();
                    sqlBeforeInsn.add(new MethodInsnNode(Opcodes.INVOKESTATIC, "java/lang/System", "currentTimeMillis", "()J",
                            false));
                    sqlBeforeInsn.add(new VarInsnNode(LSTORE, maxLocals));

                    InsnList sqlAfterInsn = new InsnList();
                    sqlAfterInsn.add(new MethodInsnNode(Opcodes.INVOKESTATIC, "java/lang/System", "currentTimeMillis", "()J",
                            false));
                    sqlAfterInsn.add(new VarInsnNode(LLOAD, maxLocals));
                    sqlAfterInsn.add(new InsnNode(LSUB));
                    sqlAfterInsn.add(new LdcInsnNode(methodPath));
                    // uuidKeyLocal是在@before创建的uuid的局部变量表位置
                    sqlAfterInsn.add(new VarInsnNode(ALOAD, uuidKeyLocal));
                    sqlAfterInsn.add(new MethodInsnNode(Opcodes.INVOKESTATIC, "TraceMonitorHandler", "sqlEnhance", "(JLjava/lang/String;Ljava/lang/String;)V",
                            false));

                    methodInsn.insert(in.getPrevious(), sqlBeforeInsn);
                    methodInsn.insert(in.getNext(), sqlAfterInsn);

                    mn.maxLocals += 1;
                    mn.maxStack += 1;
                }
            } catch (Exception e) {
                continue;
            }
        }
    }


    /**
     * 检查是否已经增强过了
     * 如果已经增强  则停止
     */
    private static boolean checkAlreadyEnhance(MethodNode mn) {
        InsnList methodInsn = mn.instructions;

        Iterator<AbstractInsnNode> j = methodInsn.iterator();
        while (j.hasNext()) {
            try {
                MethodInsnNode in = (MethodInsnNode) j.next();
                if (in.owner.contains("TraceMonitorHandler")) {
                    return true;
                }
            } catch (Exception e) {
                continue;
            }
        }
        return false;
    }

    /**
     * 前置增强 @Before
     * return ： 返回当前局部变量uuidKey位置
     */
    private static int BeforeEnhance(MethodNode mn, InsnList insns, String methodPath) {
        int maxLocals = mn.maxLocals;
        // @Before
        InsnList beforeInsn = new InsnList();
        beforeInsn.add(new MethodInsnNode(INVOKESTATIC, "java/util/UUID", "randomUUID", "()Ljava/util/UUID;",
                false));
        beforeInsn.add(new MethodInsnNode(INVOKESTATIC, "java/lang/String", "valueOf", "(Ljava/lang/Object;)Ljava/lang/String;",
                false));
        beforeInsn.add(new VarInsnNode(ASTORE, maxLocals));

        beforeInsn.add(new LdcInsnNode(methodPath));
        beforeInsn.add(new VarInsnNode(ALOAD, maxLocals));
        beforeInsn.add(new MethodInsnNode(Opcodes.INVOKESTATIC, "TraceMonitorHandler", "before", "(Ljava/lang/String;Ljava/lang/String;)V",
                false));
        insns.insert(beforeInsn);
        mn.maxLocals += 1;
        mn.maxStack += 1;
        return maxLocals;
    }

    /**
     * 后置增强 @After
     */
    private static void AfterEnhance(InsnList insns, String methodPath) {
        Iterator<AbstractInsnNode> j = insns.iterator();
        while (j.hasNext()) {
            AbstractInsnNode in = j.next();
            int op = in.getOpcode();
            if ((op >= Opcodes.IRETURN && op <= Opcodes.RETURN)) {
                InsnList il = new InsnList();
                il.add(new LdcInsnNode(methodPath));
                il.add(new MethodInsnNode(Opcodes.INVOKESTATIC, "TraceMonitorHandler", "after", "(Ljava/lang/String;)V",
                        false));
                insns.insert(in.getPrevious(), il);
            }
        }
    }

    /**
     * 异常增强 methodException
     */
    private static void TryCatchEnhance(MethodNode mn, InsnList insns, boolean isMainClass) {
        int maxLocals = mn.maxLocals;
        if (!isMainClass) {
            return;
        }
        // 异常节点Lable
        LabelNode startLabelNode = new LabelNode();
        LabelNode endLabelNode = new LabelNode();
        LabelNode exceptionHandlerLabelNode = new LabelNode();
        LabelNode returnLabelNode = new LabelNode();

        mn.tryCatchBlocks.add(new TryCatchBlockNode(startLabelNode, endLabelNode, exceptionHandlerLabelNode, "java/lang/Exception"));

        // 添加头部 try
        InsnList il = new InsnList();
        il.add(startLabelNode);
        insns.insert(il);

        // 添加尾部 catch handler
        InsnList il0 = new InsnList();
        il0.add(endLabelNode);
        il0.add(new JumpInsnNode(GOTO, returnLabelNode));

        il0.add(exceptionHandlerLabelNode);
        il0.add(new VarInsnNode(ASTORE, maxLocals));
        il0.add(new VarInsnNode(ALOAD, maxLocals));

        il0.add(new MethodInsnNode(INVOKESTATIC, "TraceMonitorHandler", "methodException", "(Ljava/lang/Exception;)V", false));
        il0.add(new VarInsnNode(ALOAD, maxLocals));

        il0.add(returnLabelNode);
        il0.add(new InsnNode(ATHROW));
        insns.add(il0);
    }
}
